from llama_cpp import Llama
import re
import json
from typing import List, Dict


class LegalAssistant:
    def __init__(self, model_path: str):
        # 初始化GGUF模型
        self.llm = Llama(
            model_path=model_path,
            n_ctx=4096,  # 上下文长度
            n_threads=8,  # CPU线程数
            verbose=False
        )
        self.dialogue_history: List[Dict] = []
        self.law_db = self.load_law_database()  # 加载法律条文数据库

    def load_law_database(self) -> Dict:
        # 示例法律数据库（实际可扩展）
        return {
            "劳动合同法": {
                38: "劳动者单方解除劳动合同的情形",
                47: "经济补偿的计算标准",
                82: "未签订书面劳动合同的赔偿"
            },
            "民法典": {
                577: "违约责任的一般规定",
                585: "违约金调整规则"
            }
        }

    def generate_response(self, user_input: str) -> str:
        # 1. 更新对话历史
        self.dialogue_history.append({"role": "user", "content": user_input})

        # 2. 构建法律增强提示词
        prompt = self.build_prompt()

        # 3. 生成回复
        response = self.llm.create_chat_completion(
            messages=prompt,
            temperature=0.3,
            max_tokens=1024
        )['choices'][0]['message']['content']

        # 4. 法律条文后处理
        processed_response = self.add_law_references(response)

        # 5. 更新历史记录
        self.dialogue_history.append({"role": "assistant", "content": processed_response})
        return processed_response

    def build_prompt(self) -> List[Dict]:
        # 系统提示词（法律领域特化）
        system_msg = {
            "role": "system",
            "content": """你是一名专业法律顾问，请遵守：
1. 必须引用中国大陆法律条文（格式：《法律名称》第XX条）
2. 分步骤给出可操作建议
3. 用【风险提示】标注潜在法律风险
4. 禁用不确定表述（如“可能”“也许”）"""
        }

        # 保留最近5轮对话
        history = self.dialogue_history[-5 * 2:]  # 每条对话含user/assistant
        return [system_msg] + history

    def add_law_references(self, text: str) -> str:
        # 自动匹配法律条文并添加解释
        pattern = r"《(.+?)》第(\d+)条"
        matches = re.finditer(pattern, text)
        for match in matches:
            law_name = match.group(1)
            article_num = int(match.group(2))
            if law_name in self.law_db and article_num in self.law_db[law_name]:
                explanation = f"（内容：{self.law_db[law_name][article_num]}）"
                text = text.replace(match.group(0), match.group(0) + explanation)
        return text

    def save_conversation(self, filename: str):
        # 保存完整对话记录
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(self.dialogue_history, f, ensure_ascii=False, indent=2)


# ----------------- 使用示例 -----------------
if __name__ == "__main__":
    # 初始化模型（需提前下载GGUF文件）
    assistant = LegalAssistant(
        model_path="deepseek-r1-distill-qwen-7b.Q4_K_M.gguf"  # 替换实际路径
    )

    # 模拟劳动纠纷咨询
    queries = [
        "公司未签劳动合同且拖欠工资怎么办？",
        "我工作已经11个月，月薪8000元",
        "如果申请劳动仲裁需要准备什么？"
    ]

    for query in queries:
        print(f"[用户] {query}")
        response = assistant.generate_response(query)
        print(f"[助理] {response}\n")

    # 保存对话记录
    assistant.save_conversation("legal_consultation.json")